#include "../Base.h"

#define PTE_BASE 0xFFFFF68000000000ull;





ULONG_PTR GetPteBase() {

	static ULONG64 pteBase = 0;
	if (pteBase) return pteBase;

	RTL_OSVERSIONINFOW version = { 0 };
	RtlGetVersion(&version);
	/*
		 Windows 10��20H2��	19042
		 Windows 10��2004)  19041
		 Windows 10��1909��	18363
		 Windows 10��1903��	18362
		 Windows 10��1809��	17763
		 Windows 10��1803��	17134
		 Windows 10��1709��	16299
		 Windows 10��1703��	15063
		 Windows 10��1607��	14393
		 Windows 10��1511��	10586
		 Windows 10	(1507)	10240

		 Windows 8.1������1��	MajorVersion = 6 MinorVersion = 3 BuildNumber = 9600
		 Windows 8.1			MajorVersion = 6 MinorVersion = 3 BuildNumber = 9200
		 Windows 8				MajorVersion = 6 MinorVersion = 2 BuildNumber = 9200
	*/
	if (version.dwBuildNumber == 7600 || version.dwBuildNumber == 7601)
	{
		pteBase = PTE_BASE;
	}
	else if (version.dwBuildNumber > 14393)
	{
		//ȡPTE
		UNICODE_STRING unName = { 0 };
		RtlInitUnicodeString(&unName, L"MmGetVirtualForPhysical");
		PUCHAR func = (PUCHAR)MmGetSystemRoutineAddress(&unName);
		pteBase = *(PULONG64)(func + 0x22);
	}
	else
	{
		pteBase = PTE_BASE;
	}

	return pteBase;
}


ULONG64 GetPte(ULONG64 VirtualAddress)
{
	ULONG64 pteBase = GetPteBase();
	return ((VirtualAddress >> 9) & 0x7FFFFFFFF8) + pteBase;
}

ULONG64 GetPde(ULONG64 VirtualAddress)
{
	ULONG64 pteBase = GetPteBase();
	ULONG64 pte = GetPte(VirtualAddress);
	return ((pte >> 9) & 0x7FFFFFFFF8) + pteBase;
}

ULONG64 GetPdpte(ULONG64 VirtualAddress)
{
	ULONG64 pteBase = GetPteBase();
	ULONG64 pde = GetPde(VirtualAddress);
	return ((pde >> 9) & 0x7FFFFFFFF8) + pteBase;
}

ULONG64 GetPml4e(ULONG64 VirtualAddress)
{
	ULONG64 pteBase = GetPteBase();
	ULONG64 ppe = GetPdpte(VirtualAddress);
	return ((ppe >> 9) & 0x7FFFFFFFF8) + pteBase;
}



BOOLEAN SetExecutePage(ULONG64 VirtualAddress, ULONG size)
{
	ULONG64 endAddress = (VirtualAddress + size) & (~0xFFF);
	ULONG64 startAddress = VirtualAddress & (~0xFFF);
	while (endAddress >= startAddress)
	{

		PPDE pde = (PPDE)GetPde(startAddress);

		if (MmIsAddressValid(pde) && pde->page_table.Present)
		{
			pde->page_table.ExecuteDisable = 0;
			pde->page_table.ReadWrite = 1;
		}


		PPTE pte = (PPTE)GetPte(startAddress);

		if (MmIsAddressValid(pte) && pte->Present)
		{
			pte->ExecuteDisable = 0;
			pte->ReadWrite = 1;
		}

		

		startAddress += PAGE_SIZE;
	}

	return TRUE;
}

PVOID AllocateMemoryNoExecute(HANDLE pid, SIZE_T size)
{
	PEPROCESS Process = NULL;
	KAPC_STATE kApcState = { 0 };
	PVOID BaseAddress = 0;
	NTSTATUS status = PsLookupProcessByProcessId(pid, &Process);

	if (!NT_SUCCESS(status))
	{
		return NULL;
	}

	if (PsGetProcessExitStatus(Process) != STATUS_PENDING)
	{
		ObDereferenceObject(Process);
		return NULL;
	}


	KeStackAttachProcess(Process, &kApcState);


	status = ZwAllocateVirtualMemory(NtCurrentProcess(), &BaseAddress, 0, &size, MEM_COMMIT, PAGE_READWRITE);

	if (NT_SUCCESS(status))
	{

		memset(BaseAddress, 0, size);
	}

	KeUnstackDetachProcess(&kApcState);

	return BaseAddress;
}

NTSTATUS FreeMemoryNoExecute(HANDLE pid, PVOID BaseAddress, SIZE_T size)
{

	PEPROCESS Process = NULL;
	KAPC_STATE kApcState = { 0 };
	NTSTATUS status = PsLookupProcessByProcessId(pid, &Process);

	if (!NT_SUCCESS(status))
	{
		return NULL;
	}

	if (PsGetProcessExitStatus(Process) != STATUS_PENDING)
	{
		ObDereferenceObject(Process);
		return NULL;
	}


	KAPC_STATE apc_state;
	KeStackAttachProcess(Process, &apc_state);

	status = ZwFreeVirtualMemory(NtCurrentProcess(), &BaseAddress, &size, MEM_RELEASE);
	if (!NT_SUCCESS(status)) {
		memset(BaseAddress, 0, size);
	}


	KeUnstackDetachProcess(&apc_state);
	ObDereferenceObject(Process);
	return status;
}

PVOID AllocateMemory(HANDLE pid, SIZE_T size) {



	PEPROCESS Process = NULL;
	KAPC_STATE kApcState = { 0 };
	PVOID BaseAddress = 0;
	NTSTATUS status = PsLookupProcessByProcessId(pid, &Process);

	if (!NT_SUCCESS(status))
	{
		return NULL;
	}

	if (PsGetProcessExitStatus(Process) != STATUS_PENDING)
	{
		ObDereferenceObject(Process);
		return NULL;
	}

	
	KAPC_STATE apc_state = {0};
	KeStackAttachProcess(Process, &apc_state);

	status = ZwAllocateVirtualMemory(NtCurrentProcess(), &BaseAddress, 0, &size, MEM_COMMIT, PAGE_READWRITE);
	
	if (NT_SUCCESS(status)) {
		memset(BaseAddress, 0, size);
		SetExecutePage((ULONG_PTR)BaseAddress, size);
	}


	KeUnstackDetachProcess(&apc_state);
	ObDereferenceObject(Process);
	return BaseAddress;
}
NTSTATUS FreeMemory(HANDLE pid, PVOID BaseAddress, SIZE_T size) {


	PEPROCESS Process = NULL;
	KAPC_STATE kApcState = { 0 };
	NTSTATUS status = PsLookupProcessByProcessId(pid, &Process);

	if (!NT_SUCCESS(status))
	{
		return NULL;
	}

	if (PsGetProcessExitStatus(Process) != STATUS_PENDING)
	{
		ObDereferenceObject(Process);
		return NULL;
	}


	KAPC_STATE apc_state;
	KeStackAttachProcess(Process, &apc_state);

	status = ZwFreeVirtualMemory(NtCurrentProcess(), &BaseAddress, &size, MEM_RELEASE);


	KeUnstackDetachProcess(&apc_state);
	ObDereferenceObject(Process);
	return status;
}
